/* Nate Deisinger (ndeisinger@wisc.edu)
 *
 * This file is part of soot-csi, a project to apply lightweight
 * instrumentation to Java programs to facilitate debugging through
 * the Soot analysis framework.
 *
 * Copyright (c) 2014, Nate Deisinger, University of Wisconsin-Madison
 *   Madison, Wisconsin 53715
 *   All Rights Reserved.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc.,
 *   51 Franklin Street, Fifth Floor,
 *   Boston, MA  02110-1301
 *   USA
 * or retrieve version 2.1 at their website:
 *   http://www.gnu.org/licenses/lgpl-2.1.html
 */
package uwmadison.jpdg;

import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import soot.SootMethod;
import soot.Value;
import soot.jimple.DefinitionStmt;
import soot.jimple.FieldRef;
import soot.jimple.internal.AbstractInstanceFieldRef;
import soot.toolkits.scalar.UnitValueBoxPair;
import edu.cwru.jpdg.graph.Edge;
import edu.cwru.jpdg.graph.ExtendedDataEdge;
import edu.cwru.jpdg.graph.ExtendedDataNode;
import edu.cwru.jpdg.graph.Graph;
import edu.cwru.jpdg.graph.Node;

/**
 * A partial implementation of an extension to the JavaPDG project to
 * allow for exporting graphs in a GraphML-friendly format.
 * 
 * Requires the JPDG classes to build.  (https://github.com/timtadh/jpdg)
 * @author Nate Deisinger
 */
public class ExtendedDataGraph extends Graph {    

    private HashMap<Integer, HashMap<Integer, List<DefinitionStmt>>> definitions;
    private HashMap<Integer, HashMap<Integer, Value>> uses;
    public HashMap<SootMethod, HashMap<Integer, Integer>> block_uids;

    public ExtendedDataGraph() {
        super();
        definitions = new HashMap<Integer, HashMap<Integer, List<DefinitionStmt>>>();
        uses = new HashMap<Integer, HashMap<Integer, Value>>();
        block_uids = new HashMap<SootMethod, HashMap<Integer, Integer>>();
    }

    public void updateDefinitions(HashMap<Integer, HashMap<Integer, List<DefinitionStmt>>> defining_stmts)
    {
        for (int i : defining_stmts.keySet())
        {
            if (!definitions.containsKey(i))
            {
                definitions.put(i, defining_stmts.get(i));
            }
            else
            {
                //System.err.println("Warning: existing node definition " + i);
            }
        }
    }

    public int addNode(String label, String extra, String package_name, String class_name, String method_name, String type, int start_l, int start_c, int end_l, int end_c) {
        int uid = next_uid;
        next_uid++;
        nodes.put(uid, new ExtendedDataNode(uid, label_num(label), extra, package_name, class_name, method_name, type, start_l, start_c, end_l, end_c, this));
        node_labels.put(uid, label_num(label));
        return uid;
    }

    public void addDataEdge(int a, int b, String type, UnitValueBoxPair u)
    {
        super.addEdge(a, b, type);
        if (!uses.containsKey(b))
        {
            uses.put(b, new  HashMap<Integer, Value>());
        }
        HashMap<Integer, Value> bSet = uses.get(b);
        //int var = d_stmt.getRightOp().
        soot.Local var = null;
        try {
            var  = (soot.Local)u.getValueBox().getValue();
        } catch (java.lang.ClassCastException e) {
            // It will be a field reference if it is not a local.
            // In the future we want to dataflow through field
            // references however that will always be handled by
            // another method as it will be inter-procedural.
            System.err.println("RightOp was not a local");
            System.err.println(u.getUnit());
            return;
        }
        //System.out.println(var);
        if (!bSet.containsKey(var.getNumber()))
        {
            bSet.put(var.getNumber(), var);
        }
    }

    public String Serialize()
    {
        //String start = super.Serialize();
        StringBuilder sb = new StringBuilder();   
        for (Node n : nodes.values()) {
            //Pull out must-def, may-def, and may-use sets
            //System.out.println(definitions.toString());
            HashMap<Integer, List<DefinitionStmt>> defs = definitions.get(n.id);

            HashSet<String> mustDef = new HashSet<String>();
            HashSet<String> mayDef = new HashSet<String>();
            if (defs != null)
            {
                for (List<DefinitionStmt> defSet : defs.values())
                {
                    for (DefinitionStmt s : defSet)
                    {
                        if (s.getLeftOp() instanceof AbstractInstanceFieldRef)
                        {
                            AbstractInstanceFieldRef ourRef = (AbstractInstanceFieldRef) s.getLeftOp();
                            // mayDef.add(ourRef.getBase().toString());           
                            mayDef.add(ourRef.toString());
                        }
                        else if (s.getLeftOp().getType() instanceof soot.PrimType)
                        {
                            mustDef.add(s.getLeftOp().toString());
                        }
                        else
                        {
                            mustDef.add(s.getLeftOp().toString());
                            //mayDef.add(s.getLeftOp().toString());
                        }
                    }
                }
            }
            HashMap<Integer, Value> nodeUses = uses.get(n.id);
            HashSet<String> mayUse = new HashSet<String>();
            if (nodeUses != null)
            {
                for (Value v : nodeUses.values())
                {
                    if (v instanceof FieldRef)
                    {
                        FieldRef ourRef = (FieldRef)v;
                        mayUse.add(ourRef.getField().getSignature());
                    }
                    else
                    {
                        mayUse.add(v.toString());
                    }
                }
            }
            sb.append(((ExtendedDataNode)n).dataSerialize(mustDef, mayDef, mayUse));
            sb.append("\n");
        }    
        for (Map.Entry<String,HashMap<Integer,Set<Integer>>> E : edges.entrySet()) {
            String e_label = E.getKey();
            HashMap<Integer,Set<Integer>> edges = E.getValue();
            for (Map.Entry<Integer,Set<Integer>> e : edges.entrySet()) {
                int i = e.getKey();
                for (int j : e.getValue()) {
                    sb.append((new ExtendedDataEdge(i, j, e_label, this)).Serialize());
                    sb.append("\n");
                }
            }
        }
        return sb.toString();
    }

}